import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import cv2
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont
import random
import time
from typing import Dict, List, Tuple, Optional

class EyeDetector(nn.Module):
    def __init__(self):
        super(EyeDetector, self).__init__()
        # 使用预训练的MobileNetV2作为基础网络
        self.base_model = torchvision.models.mobilenet_v2(pretrained=True)
        # 移除分类层
        self.features = self.base_model.features
        
        # 自定义眼睛检测头
        self.eye_detector = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(1280, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 4)  # 输出为[left_eye_x, left_eye_y, right_eye_x, right_eye_y]
        )
        
        # 瞳孔检测头
        self.pupil_detector = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(1280, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 4)  # 输出为[left_pupil_x, left_pupil_y, right_pupil_x, right_pupil_y]
        )
        
    def forward(self, x):
        # 提取特征
        features = self.features(x)
        
        # 检测眼睛位置
        eye_coords = self.eye_detector(features)
        
        # 检测瞳孔位置
        pupil_coords = self.pupil_detector(features)
        
        return eye_coords, pupil_coords

    def extract_eye_regions(self, image, eye_coords):
        """从原始图像中提取眼睛区域"""
        left_eye_x, left_eye_y, right_eye_x, right_eye_y = eye_coords.cpu().numpy()[0]
        
        # 计算眼睛区域（添加边距）
        eye_width = 60
        eye_height = 40
        
        left_eye_region = image[
            max(0, int(left_eye_y - eye_height/2)):min(image.shape[0], int(left_eye_y + eye_height/2)),
            max(0, int(left_eye_x - eye_width/2)):min(image.shape[1], int(left_eye_x + eye_width/2))
        ]
        
        right_eye_region = image[
            max(0, int(right_eye_y - eye_height/2)):min(image.shape[0], int(right_eye_y + eye_height/2)),
            max(0, int(right_eye_x - eye_width/2)):min(image.shape[1], int(right_eye_x + eye_width/2))
        ]
        
        return left_eye_region, right_eye_region
    
class VisionChartGenerator:
    def __init__(self, chart_type="snellen"):
        self.chart_type = chart_type
        self.chart_size = (800, 600)  # 默认图表大小
        self.background_color = (255, 255, 255)  # 白色背景
        self.text_color = (0, 0, 0)  # 黑色文本
        
        # 不同视力等级对应的字符大小（像素）
        self.vision_sizes = {
            "20/200": 72,  # 0.1视力
            "20/100": 36,  # 0.2视力
            "20/70": 24,   # 0.3视力
            "20/50": 18,   # 0.4视力
            "20/40": 14,   # 0.5视力
            "20/30": 10,   # 0.6-0.7视力
            "20/25": 8,    # 0.8视力
            "20/20": 7,    # 1.0视力
            "20/15": 5,    # 1.2视力
            "20/10": 4,    # 1.5视力
        }
        
        # 不同视力对应的近视度数（大致值，仅供参考）
        self.vision_to_diopter = {
            "20/200": -6.00,
            "20/100": -3.00,
            "20/70": -2.00,
            "20/50": -1.50,
            "20/40": -1.00,
            "20/30": -0.75,
            "20/25": -0.50,
            "20/20": 0.00,
            "20/15": 0.00,
            "20/10": 0.00,
        }
        
        # 斯内伦视力表字符集
        self.snellen_chars = ['E', 'F', 'P', 'T', 'O', 'Z', 'L', 'C', 'D']
        # Landolt C字符（方向）
        self.landolt_directions = [0, 45, 90, 135, 180, 225, 270, 315]
        
    def generate_chart(self, vision_level="20/40", num_chars=5):
        """生成指定视力级别的视力表"""
        if vision_level not in self.vision_sizes:
            vision_level = "20/40"  # 默认视力级别
            
        # 创建空白画布
        chart = Image.new('RGB', self.chart_size, self.background_color)
        draw = ImageDraw.Draw(chart)
        
        # 字符大小
        char_size = self.vision_sizes[vision_level]
        
        if self.chart_type == "snellen":
            # 随机选择字符
            selected_chars = random.choices(self.snellen_chars, k=num_chars)
            
            # 加载适当大小的字体
            try:
                font = ImageFont.truetype("arial.ttf", char_size)
            except IOError:
                font = ImageFont.load_default()
            
            # 计算开始位置，使字符居中
            total_width = num_chars * char_size * 1.2  # 估计总宽度
            start_x = (self.chart_size[0] - total_width) // 2
            y_position = (self.chart_size[1] - char_size) // 2
            
            # 绘制字符
            for i, char in enumerate(selected_chars):
                x_position = start_x + i * char_size * 1.2
                draw.text((x_position, y_position), char, font=font, fill=self.text_color)
                
            return chart, selected_chars
            
        elif self.chart_type == "landolt":
            # 生成Landolt C图表（使用简单绘制替代文字）
            c_size = char_size
            gap_size = c_size // 5
            directions = random.choices(self.landolt_directions, k=num_chars)
            
            # 计算开始位置，使字符居中
            total_width = num_chars * c_size * 1.5
            start_x = (self.chart_size[0] - total_width) // 2
            y_position = (self.chart_size[1] - c_size) // 2
            
            landolt_cs = []
            for i, angle in enumerate(directions):
                x_position = start_x + i * c_size * 1.5
                
                # 绘制Landolt C
                draw.ellipse(
                    [x_position, y_position, x_position + c_size, y_position + c_size], 
                    outline=self.text_color, width=max(1, c_size//10)
                )
                
                # 根据角度计算缺口位置
                gap_x = x_position + c_size/2 + (c_size/2 - gap_size/2) * np.cos(np.radians(angle))
                gap_y = y_position + c_size/2 + (c_size/2 - gap_size/2) * np.sin(np.radians(angle))
                
                # 绘制缺口（用白色覆盖一部分圆）
                draw.ellipse(
                    [gap_x - gap_size/2, gap_y - gap_size/2, gap_x + gap_size/2, gap_y + gap_size/2], 
                    fill=self.background_color, outline=self.background_color
                )
                
                landolt_cs.append(int(angle // 45))  # 将角度转换为0-7的数字表示方向
                
            return chart, landolt_cs
        
        else:
            raise ValueError(f"不支持的视力表类型: {self.chart_type}")
    
    def simulate_blurry_vision(self, chart, diopter):
        """模拟不同度数的近视效果"""
        # 将度数转换为模糊内核大小 (简化模型)
        if diopter >= 0:
            return chart  # 无需模糊
            
        blur_amount = int(abs(diopter) * 3)  # 度数越高，模糊越严重
        blur_amount = max(1, min(15, blur_amount))  # 限制在合理范围
        
        # 转换为numpy数组进行处理
        chart_np = np.array(chart)
        
        # 应用高斯模糊
        blurred = cv2.GaussianBlur(chart_np, (blur_amount*2+1, blur_amount*2+1), 0)
        
        return Image.fromarray(blurred)
    

class GazeEstimator(nn.Module):
    def __init__(self):
        super(GazeEstimator, self).__init__()
        
        # 使用ResNet-18作为基础
        self.base_model = torchvision.models.resnet18(pretrained=True)
        num_ftrs = self.base_model.fc.in_features
        
        # 替换最后的全连接层
        self.base_model.fc = nn.Identity()
        
        # 添加注视方向估计器
        self.gaze_regressor = nn.Sequential(
            nn.Linear(num_ftrs, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 2)  # [gaze_x, gaze_y] - 归一化坐标
        )
        
    def forward(self, left_eye, right_eye):
        # 合并两只眼睛的特征
        left_features = self.base_model(left_eye)
        right_features = self.base_model(right_eye)
        
        # 计算注视向量
        left_gaze = self.gaze_regressor(left_features)
        right_gaze = self.gaze_regressor(right_features)
        
        # 平均两只眼睛的注视方向
        gaze = (left_gaze + right_gaze) / 2
        
        return gaze